import json
import os
from typing import Dict, List, Any, Optional
from collections import deque, defaultdict
import sys
import time
from typing import Dict, List, Any, Optional
sys.stdout.reconfigure(encoding='utf-8')
os.system('chcp 65001 >nul')
class PDCGAnalyzer:
    """
    参数化数据流调用图 (PDCG) 分析器
    基于AST结构生成JavaScript代码的调用图和数据流分析
    """
    
    def __init__(self):
        self.nodes = []
        self.edges = []
        self.node_id_counter = 0
        self.ast_cache = None
        
        # 节点映射
        self.ast_to_node = {}  # AST节点 -> 生成的节点ID
        
        # 处理队列
        self.processing_queue = deque()
        
        # 新增：函数定义符号表
        self.function_definitions = {}  # 函数名 -> 定义节点ID
    def analyze_pdcg_from_ast(self, ast_file_path: str) -> Dict[str, Any]:
        """主分析流程"""
        print(f"开始PDCG分析: {ast_file_path}")
        
        # 1. 加载AST
        if not self._load_ast(ast_file_path):
            return {"error": "Failed to load AST"}
        
        # 2. 创建根节点并初始化队列
        self._create_root_and_enqueue()
        
        # 3. 处理队列，生成所有节点和边
        self._process_queue()
        
        # 4. 验证和统计
        result = self._validate_and_generate_stats()
        
        print(f"PDCG分析完成: {len(self.nodes)}个节点, {len(self.edges)}条边")
        return result
    
    def _load_ast(self, ast_file_path: str) -> bool:
        """加载AST文件"""
        try:
            with open(ast_file_path, 'r', encoding='utf-8') as f:
                ast_data = json.load(f)
            
            if not ast_data.get('success', False):
                print("AST解析失败")
                return False
            
            self.ast_cache = ast_data['ast']['program']
            print(f"AST加载成功，程序体包含 {len(self.ast_cache.get('body', []))} 个顶级语句")
            return True
            
        except Exception as e:
            print(f"加载AST失败: {e}")
            return False
    
    def _create_root_and_enqueue(self):
        """创建根节点并初始化队列"""
        # 创建文件根节点
        file_id = self._create_node('FILE', 'file')
        
        # 将程序体加入队列
        program_body = self.ast_cache.get('body', [])
        for stmt in program_body:
            if stmt is not None:  # 添加None检查
                self.processing_queue.append({
                    'ast_node': stmt,
                    'parent_id': file_id,
                    'context': 'top_level'
                })
        
        print(f"根节点创建完成，队列初始化：{len(self.processing_queue)} 个待处理项")
    def _process_queue(self):
        """处理队列"""
        level = 1
        
        while self.processing_queue:
            current_level_size = len(self.processing_queue)
            print(f"处理第 {level} 层，包含 {current_level_size} 个节点")
            
            for _ in range(current_level_size):
                if not self.processing_queue:
                    break
                
                item = self.processing_queue.popleft()
                self._process_queue_item(item)
            
            level += 1
            if level > 50:
                print(f"达到最大层数限制，停止处理")
                break
        
        print(f"队列处理完成，总共处理了 {level-1} 层")
    
    def _process_queue_item(self, item: Dict):
        """处理队列项目"""
        # 添加None检查
        if not item or not isinstance(item, dict):
            return
            
        ast_node = item.get('ast_node')
        parent_id = item.get('parent_id')
        context = item.get('context')
        
        # 确保ast_node不为None且是字典
        if not ast_node or not isinstance(ast_node, dict):
            return
            
        node_type = ast_node.get('type')
        if not node_type:
            return
        
        # 根据节点类型分发处理
        if node_type == 'FunctionDeclaration':
            self._handle_function_declaration(ast_node, parent_id, context)
        elif node_type == 'ClassDeclaration':
            self._handle_class_declaration(ast_node, parent_id, context)
        elif node_type == 'CallExpression':
            self._handle_call_expression(ast_node, parent_id, context)
        elif node_type in ['FunctionExpression', 'ArrowFunctionExpression']:
            self._handle_function_expression(ast_node, parent_id, context)
        elif node_type == 'NewExpression':
            self._handle_new_expression(ast_node, parent_id, context)
        elif node_type == 'VariableDeclaration':
            self._handle_variable_declaration(ast_node, parent_id, context)
        elif node_type == 'ExpressionStatement':
            self._handle_expression_statement(ast_node, parent_id, context)
        elif node_type == 'AssignmentExpression':
        # 新增：递归处理赋值左右两边
            left = ast_node.get('left')
            right = ast_node.get('right')
            if left:
                self._enqueue_child(left, parent_id, 'assignment_left')
            if right:
                self._enqueue_child(right, parent_id, 'assignment_right')
        elif node_type == 'MemberExpression':
        # 新增：递归处理对象和属性
            obj = ast_node.get('object')
            prop = ast_node.get('property')
            if obj:
                self._enqueue_child(obj, parent_id, 'member_object')
            if prop:
                self._enqueue_child(prop, parent_id, 'member_property')
        elif node_type in ['IfStatement', 'ForStatement', 'WhileStatement', 'BlockStatement']:
            self._handle_control_flow(ast_node, parent_id, context)
        else:
            self._handle_generic_node(ast_node, parent_id, context)

    def _handle_call_expression(self, ast_node: Dict, parent_id: str, context: str):
        """处理CallExpression - 核心方法"""
        call_analysis = self._analyze_call_structure(ast_node)
        call_id = self._create_node(
            'CALL',
            'call',
            callee_name=call_analysis['name'],
            is_chained=call_analysis['is_chained'],
            is_member_call=call_analysis['is_member'],
            call_depth=call_analysis['call_depth'],
            chain_position=call_analysis['chain_position']
        )
        self.ast_to_node[id(ast_node)] = call_id
    
        # 建立父子关系（原逻辑不变）
        if call_analysis['is_chained']:
            base_call_ast = call_analysis.get('base_call_ast')
            if base_call_ast:
                base_call_id = self._ensure_call_processed(base_call_ast, parent_id, context)
                if base_call_id:
                    self._add_edge(base_call_id, call_id, 'chained_call')
                else:
                    edge_type = 'calls' if context in ['top_level', 'function_body'] else 'nested_in'
                    self._add_edge(parent_id, call_id, edge_type)
            else:
                edge_type = 'calls' if context in ['top_level', 'function_body'] else 'nested_in'
                self._add_edge(parent_id, call_id, edge_type)
        else:
            edge_type = 'calls' if context in ['top_level', 'function_body'] else 'nested_in'
            self._add_edge(parent_id, call_id, edge_type)
    
        # 新增：检查并关联自定义函数调用
        callee = ast_node.get('callee', {})
        if callee.get('type') == 'Identifier':
            func_name_to_check = callee.get('name')
            if func_name_to_check and func_name_to_check in self.function_definitions:
                definition_node_id = self.function_definitions[func_name_to_check]
                self._add_edge(call_id, definition_node_id, 'resolves_to_def')
                #print(f"  自定义函数调用已关联: {func_name_to_check}() -> {definition_node_id}")
    
        # 处理参数
        self._handle_call_arguments(ast_node.get('arguments', []), call_id)
    
        # 修复点1：递归处理callee为CallExpression
        if callee:
            if callee.get('type') in ['FunctionExpression', 'ArrowFunctionExpression']:
                body = callee.get('body')
                if body:
                    self._enqueue_function_body(body, call_id)
            elif callee.get('type') == 'CallExpression':
                self._enqueue_child(callee, call_id, 'callee')
            elif callee.get('type') == 'MemberExpression':
                self._enqueue_child(callee, call_id, 'callee')
    def _ensure_call_processed(self, call_ast: Dict, parent_id: str, context: str) -> Optional[str]:
        """确保调用已被处理，如果没有则立即处理"""
        # 检查是否已处理
        call_id = self.ast_to_node.get(id(call_ast))
        if call_id:
            return call_id
        
        # 如果没有处理，立即处理
        if call_ast.get('type') == 'CallExpression':
            self._handle_call_expression(call_ast, parent_id, context)
            return self.ast_to_node.get(id(call_ast))
        
        return None
    
    def _analyze_call_structure(self, ast_node: Dict) -> Dict:
        """分析调用结构"""
        if ast_node.get('type') != 'CallExpression':
            return {
                'is_chained': False, 
                'name': 'unknown', 
                'display_name': 'unknown',
                'is_member': False,
                'call_depth': 0,
                'chain_position': 'unknown'
            }
        
        callee = ast_node.get('callee', {})
        
        if callee.get('type') == 'Identifier':
            # 简单调用：func()
            name = callee.get('name', 'unknown')
            return {
                'name': name,
                'display_name': name,
                'is_chained': False,
                'is_member': False,
                'call_depth': 0,
                'chain_position': 'simple_call'
            }
        
        elif callee.get('type') == 'MemberExpression':
            obj = callee.get('object', {})
            prop = callee.get('property', {})
            prop_name = prop.get('name', 'unknown')
            
            if obj.get('type') == 'CallExpression':
                # 链式调用：obj().method()
                base_analysis = self._analyze_call_chain_depth(obj)
                
                return {
                    'name': f"{base_analysis['name']}.{prop_name}",
                    'display_name': f"(...).{prop_name}",
                    'is_chained': True,
                    'is_member': True,
                    'call_depth': base_analysis['call_depth'] + 1,
                    'chain_position': 'chained',
                    'base_call_ast': obj
                }
            else:
                # 普通成员调用：obj.method()
                obj_name = self._extract_expression_name(obj)
                return {
                    'name': f"{obj_name}.{prop_name}",
                    'display_name': f"{obj_name}.{prop_name}",
                    'is_chained': False,
                    'is_member': True,
                    'call_depth': 0,
                    'chain_position': 'member_call'
                }
        
        return {
            'name': 'unknown',
            'display_name': 'unknown',
            'is_chained': False,
            'is_member': False,
            'call_depth': 0,
            'chain_position': 'unknown'
        }
    
    def _analyze_call_chain_depth(self, call_ast: Dict) -> Dict:
        """分析调用链的深度"""
        if call_ast.get('type') != 'CallExpression':
            return {'call_depth': 0, 'name': 'unknown'}
        
        callee = call_ast.get('callee', {})
        
        if callee.get('type') == 'Identifier':
            # 基础调用
            return {
                'call_depth': 0,
                'name': callee.get('name', 'unknown')
            }
        
        elif callee.get('type') == 'MemberExpression':
            obj = callee.get('object', {})
            prop = callee.get('property', {})
            prop_name = prop.get('name', 'unknown')
            
            if obj.get('type') == 'CallExpression':
                # 递归分析更深的链
                deeper_analysis = self._analyze_call_chain_depth(obj)
                return {
                    'call_depth': deeper_analysis['call_depth'] + 1,
                    'name': f"{deeper_analysis['name']}.{prop_name}"
                }
            else:
                # 成员调用作为基础
                obj_name = self._extract_expression_name(obj)
                return {
                    'call_depth': 0,
                    'name': f"{obj_name}.{prop_name}"
                }
        
        return {'call_depth': 0, 'name': 'unknown'}
    
    def _handle_call_arguments(self, args: List[Dict], call_id: str):
        """处理调用参数 - 最终简化版"""
        for i, arg in enumerate(args):
            if arg is None:  # 添加None检查
                # 创建null参数节点
                arg_id = self._create_node(
                    'ARGUMENT',
                    'null',
                    arg_index=i,
                    content='null'
                )
                self._add_edge(call_id, arg_id, 'has_arg')
                continue
                
            # 其余逻辑完全不变
            arg_content = self._extract_argument_content(arg)
            
            arg_id = self._create_node(
                'ARGUMENT',
                arg_content['type'],
                arg_index=i,
                content=arg_content['content']
            )
            
            self._add_edge(call_id, arg_id, 'has_arg')
            self._enqueue_child(arg, arg_id, 'call_argument')
    def _extract_argument_content(self, arg: Dict) -> Dict:
        """提取参数内容 - 最终版：保持原始内容，简化返回结构"""
        if not arg:
            return {'type': 'null', 'content': 'null'}
        
        arg_type = arg.get('type')
        
        # 字面量类型 - 按原代码保持
        if arg_type in ['Literal', 'StringLiteral', 'NumericLiteral', 'BooleanLiteral']:
            value = arg.get('value')
            if isinstance(value, str):
                return {'type': 'string', 'content': value}
            elif isinstance(value, (int, float)):
                return {'type': 'number', 'content': str(value)}
            elif isinstance(value, bool):
                # 保持JavaScript原始boolean值
                return {'type': 'boolean', 'content': 'true' if value else 'false'}
            else:
                content = str(value) if value is not None else 'null'
                return {'type': 'literal', 'content': content}
        
        # 标识符
        elif arg_type == 'Identifier':
            name = arg.get('name', 'unknown')
            return {'type': 'variable', 'content': name}
        
        # 模板字符串 - 完整拼接，不做任何简化
        elif arg_type == 'TemplateLiteral':
            return self._extract_template_literal_content(arg)
        
        # 函数调用
        elif arg_type == 'CallExpression':
            call_analysis = self._analyze_call_structure(arg)
            return {'type': 'function_call', 'content': call_analysis['name']}
        
        # 对象表达式 - 完整拼接
        elif arg_type == 'ObjectExpression':
            return self._extract_object_expression_content(arg)
        
        # 数组表达式 - 完整拼接
        elif arg_type == 'ArrayExpression':
            return self._extract_array_expression_content(arg)
        
        # 成员访问
        elif arg_type == 'MemberExpression':
            return self._extract_member_expression_content(arg)
        
        # 二元表达式
        elif arg_type == 'BinaryExpression':
            return self._extract_binary_expression_content(arg)
        
        # 更新表达式
        elif arg_type == 'UpdateExpression':
            return self._extract_update_expression_content(arg)
        
        # 正则表达式
        elif arg_type == 'RegExpLiteral':
            pattern = arg.get('pattern', '')
            flags = arg.get('flags', '')
            return {'type': 'regex', 'content': f"/{pattern}/{flags}"}
        
        # 其他类型
        else:
            return {'type': 'unknown', 'content': arg_type}
    
    def _extract_template_literal_content(self, template: Dict) -> Dict:
        """提取模板字符串内容 - 完整保留原始内容"""
        quasis = template.get('quasis', [])
        expressions = template.get('expressions', [])
        
        # 完整拼接模板字符串，不做任何简化
        template_parts = []
        expr_index = 0
        
        for i, quasi in enumerate(quasis):
            raw_value = quasi.get('value', {}).get('raw', '')
            template_parts.append(raw_value)
            
            if expr_index < len(expressions):
                expr = expressions[expr_index]
                expr_content = self._extract_expression_for_template(expr)
                template_parts.append(f"${{{expr_content}}}")
                expr_index += 1
        
        full_content = ''.join(template_parts)
        
        return {'type': 'template', 'content': full_content}
    
    def _extract_expression_for_template(self, expr: Dict) -> str:
        """为模板字符串提取表达式内容 - 直接拼接"""
        if expr is None:  # 添加None检查
            return 'unknown'
            
        expr_type = expr.get('type')
        
        if expr_type == 'Identifier':
            return expr.get('name', 'unknown')
        elif expr_type == 'UpdateExpression':
            operator = expr.get('operator', '')
            argument_node = expr.get('argument', {})
            if argument_node is not None:  # 添加None检查
                argument = self._extract_expression_for_template(argument_node)
            else:
                argument = 'unknown'
            prefix = expr.get('prefix', True)
            return f"{operator}{argument}" if prefix else f"{argument}{operator}"
        elif expr_type in ['Literal', 'StringLiteral', 'NumericLiteral']:
            value = expr.get('value')
            return f'"{value}"' if isinstance(value, str) else str(value)
        elif expr_type == 'BooleanLiteral':
            value = expr.get('value')
            return 'true' if value else 'false'
        elif expr_type == 'CallExpression':
            call_analysis = self._analyze_call_structure(expr)
            args = expr.get('arguments', [])
            arg_contents = []
            for arg in args:
                if arg is not None:  # 添加None检查
                    arg_contents.append(self._extract_expression_for_template(arg))
            return f"{call_analysis['name']}({', '.join(arg_contents)})"
        elif expr_type == 'MemberExpression':
            obj_node = expr.get('object', {})
            prop_node = expr.get('property', {})
            
            if obj_node is not None:  # 添加None检查
                obj = self._extract_expression_for_template(obj_node)
            else:
                obj = 'unknown'
                
            if prop_node is not None:  # 添加None检查
                prop = prop_node.get('name', 'unknown')
            else:
                prop = 'unknown'
                
            computed = expr.get('computed', False)
            return f"{obj}[{prop}]" if computed else f"{obj}.{prop}"
        elif expr_type == 'BinaryExpression':
            left_node = expr.get('left', {})
            right_node = expr.get('right', {})
            
            if left_node is not None:  # 添加None检查
                left = self._extract_expression_for_template(left_node)
            else:
                left = 'unknown'
                
            if right_node is not None:  # 添加None检查
                right = self._extract_expression_for_template(right_node)
            else:
                right = 'unknown'
                
            operator = expr.get('operator', '')
            return f"{left} {operator} {right}"
        else:
            return expr_type if expr_type else 'unknown'
    def _extract_object_expression_content(self, obj_expr: Dict) -> Dict:
        """提取对象表达式内容 - 完整拼接"""
        properties = obj_expr.get('properties', [])
        
        if not properties:
            return {'type': 'object', 'content': '{}'}
        
        # 完整拼接所有属性
        prop_strings = []
        for prop in properties:
            if prop.get('type') in ['Property', 'ObjectProperty']:
                key = self._extract_property_key(prop)
                value = self._extract_property_value(prop)
                prop_strings.append(f"{key}: {value}")
        
        obj_content = "{" + ", ".join(prop_strings) + "}"
        return {'type': 'object', 'content': obj_content}
    
    def _extract_property_key(self, prop: Dict) -> str:
        """提取对象属性键 - 直接返回字符串"""
        key = prop.get('key', {})
        key_type = key.get('type')
        
        if key_type == 'Identifier':
            return key.get('name', 'unknown')
        elif key_type in ['Literal', 'StringLiteral']:
            value = key.get('value')
            return f'"{value}"' if isinstance(value, str) else str(value)
        else:
            return 'unknown'
    
    def _extract_property_value(self, prop: Dict) -> str:
        """提取对象属性值 - 保持原始boolean值"""
        value = prop.get('value', {})
        
        # 特殊处理boolean值
        if value.get('type') == 'BooleanLiteral':
            bool_value = value.get('value')
            return 'true' if bool_value else 'false'
        
        value_content = self._extract_argument_content(value)
        return value_content['content']
    
    def _extract_array_expression_content(self, arr_expr: Dict) -> Dict:
        """提取数组表达式内容 - 完整内容提取"""
        elements = arr_expr.get('elements', [])
        
        if not elements:
            return {'type': 'array', 'content': '[]'}
        
        # 递归处理每个元素，获取具体内容
        elem_strings = []
        for elem in elements:
            if elem is None:
                elem_strings.append('null')
            else:
                # 递归调用，获取元素的具体内容
                elem_content = self._extract_argument_content(elem)
                elem_strings.append(elem_content['content'])
        
        # 拼接完整数组内容
        arr_content = "[" + ", ".join(elem_strings) + "]"
        return {'type': 'array', 'content': arr_content}
    def _extract_member_expression_content(self, member_expr: Dict) -> Dict:
        """提取成员表达式内容 - 完整拼接"""
        obj = member_expr.get('object', {})
        prop = member_expr.get('property', {})
        computed = member_expr.get('computed', False)
        
        obj_content = self._extract_argument_content(obj)
        prop_content = self._extract_argument_content(prop)
        
        if computed:
            full_content = f"{obj_content['content']}[{prop_content['content']}]"
        else:
            full_content = f"{obj_content['content']}.{prop_content['content']}"
        
        return {'type': 'member_access', 'content': full_content}
    
    def _extract_binary_expression_content(self, binary_expr: Dict) -> Dict:
        """提取二元表达式内容 - 完整拼接"""
        left = binary_expr.get('left', {})
        right = binary_expr.get('right', {})
        operator = binary_expr.get('operator', '?')
        
        left_content = self._extract_argument_content(left)
        right_content = self._extract_argument_content(right)
        
        full_content = f"{left_content['content']} {operator} {right_content['content']}"
        return {'type': 'binary_expression', 'content': full_content}
    
    def _extract_update_expression_content(self, update_expr: Dict) -> Dict:
        """提取更新表达式内容 - 完整拼接"""
        argument = update_expr.get('argument', {})
        operator = update_expr.get('operator', '++')
        prefix = update_expr.get('prefix', True)
        
        arg_content = self._extract_argument_content(argument)
        
        if prefix:
            full_content = f"{operator}{arg_content['content']}"
        else:
            full_content = f"{arg_content['content']}{operator}"
        
        return {'type': 'update_expression', 'content': full_content}
    
    def _handle_function_declaration(self, ast_node: Dict, parent_id: str, context: str):
        """处理函数声明"""
        # 修复这里：正确处理id为None的情况
        id_node = ast_node.get('id')
        if id_node is not None:
            func_name = id_node.get('name', 'anonymous')
        else:
            func_name = 'anonymous'
        
        func_id = self._create_node(
            'FUNCTION_DEF',
            'function_def',
            function_name=func_name,
            is_async=ast_node.get('async', False),
            is_generator=ast_node.get('generator', False)
        )
        
        # 新增：记录函数定义到符号表
        if func_name != 'anonymous':
            self.function_definitions[func_name] = func_id
            print(f"  函数定义已记录: {func_name} -> {func_id}")
        
        self._add_edge(parent_id, func_id, 'defines')
        self._handle_function_parameters(ast_node.get('params', []), func_id)
        
        body = ast_node.get('body')
        if body is not None:  # 添加None检查
            self._enqueue_function_body(body, func_id)
    def _handle_class_declaration(self, ast_node: Dict, parent_id: str, context: str):
        """处理类声明"""
        # 添加None检查
        id_node = ast_node.get('id')
        if id_node is not None:
            class_name = id_node.get('name', 'anonymous')
        else:
            class_name = 'anonymous'
        
        class_id = self._create_node(
            'CLASS',
            'class',
            class_name=class_name
        )
        
        self._add_edge(parent_id, class_id, 'defines')
        
        class_body = ast_node.get('body', {})
        if class_body and class_body.get('type') == 'ClassBody':
            methods = class_body.get('body', [])
            for method in methods:
                if method is not None:  # 添加None检查
                    self._enqueue_child(method, class_id, 'class_member')
    def _handle_function_expression(self, ast_node: Dict, parent_id: str, context: str):
        """处理函数表达式"""
        is_callback = context in ['call_argument', 'array_element', 'object_property']
        
        func_id = self._create_node(
            'ANONYMOUS_FUNCTION',
            'anonymous_function',
            is_callback=is_callback,
            is_async=ast_node.get('async', False)
        )
        
        edge_type = 'callback' if is_callback else 'defines'
        self._add_edge(parent_id, func_id, edge_type)
        
        self._handle_function_parameters(ast_node.get('params', []), func_id)
        
        body = ast_node.get('body')
        if body:
            self._enqueue_function_body(body, func_id)
    
    def _handle_new_expression(self, ast_node: Dict, parent_id: str, context: str):
        """处理new表达式"""
        constructor_name = self._extract_constructor_name(ast_node)
        
        call_id = self._create_node(
            'CALL',
            'call',
            callee_name=f"new {constructor_name}",
            is_constructor=True
        )
        
        edge_type = 'calls' if context in ['top_level', 'function_body'] else 'nested_in'
        self._add_edge(parent_id, call_id, edge_type)
        
        self._handle_call_arguments(ast_node.get('arguments', []), call_id)
    
    def _handle_variable_declaration(self, ast_node: Dict, parent_id: str, context: str):
        """处理变量声明"""
        for declarator in ast_node.get('declarations', []):
            if declarator is not None and declarator.get('init') is not None:  # 添加None检查
                self._enqueue_child(declarator['init'], parent_id, 'variable_init')
    def _handle_expression_statement(self, ast_node: Dict, parent_id: str, context: str):
        """处理表达式语句"""
        expression = ast_node.get('expression')
        if expression is not None:  # 添加None检查
            self._enqueue_child(expression, parent_id, context)
    def _handle_control_flow(self, ast_node: Dict, parent_id: str, context: str):
        """处理控制流语句"""
        node_type = ast_node.get('type')
        
        if node_type == 'IfStatement':
            test_node = ast_node.get('test')
            if test_node is not None:  # 添加None检查
                self._enqueue_child(test_node, parent_id, 'condition')
                
            consequent_node = ast_node.get('consequent')
            if consequent_node is not None:  # 添加None检查
                self._enqueue_child(consequent_node, parent_id, 'branch')
                
            alternate_node = ast_node.get('alternate')
            if alternate_node is not None:  # 添加None检查
                self._enqueue_child(alternate_node, parent_id, 'branch')
        
        elif node_type == 'ForStatement':
            for field in ['init', 'test', 'update', 'body']:
                field_node = ast_node.get(field)
                if field_node is not None:  # 添加None检查
                    context_map = {
                        'init': 'loop_init',
                        'test': 'loop_condition', 
                        'update': 'loop_update',
                        'body': 'loop_body'
                    }
                    self._enqueue_child(field_node, parent_id, context_map[field])
        
        elif node_type == 'WhileStatement':
            test_node = ast_node.get('test')
            if test_node is not None:  # 添加None检查
                self._enqueue_child(test_node, parent_id, 'loop_condition')
                
            body_node = ast_node.get('body')
            if body_node is not None:  # 添加None检查
                self._enqueue_child(body_node, parent_id, 'loop_body')
        
        elif node_type == 'BlockStatement':
            body_nodes = ast_node.get('body', [])
            for stmt in body_nodes:
                if stmt is not None:  # 添加None检查
                    self._enqueue_child(stmt, parent_id, 'function_body')
    def _handle_generic_node(self, ast_node: Dict, parent_id: str, context: str):
        """处理通用节点"""
        if ast_node is None:  # 添加None检查
            return
            
        for key, value in ast_node.items():
            if key in ['loc', 'start', 'end', 'range', 'comments']:
                continue
            
            if isinstance(value, list):
                for item in value:
                    if item is not None and isinstance(item, dict) and item.get('type'):  # 添加None检查
                        self._enqueue_child(item, parent_id, context)
            
            elif value is not None and isinstance(value, dict) and value.get('type'):  # 添加None检查
                self._enqueue_child(value, parent_id, context)
    def _handle_function_parameters(self, params: List[Dict], func_id: str):
        """处理函数参数"""
        if not params:  # 添加None检查
            return
            
        for param in params:
            if param is not None:  # 添加None检查
                param_name = self._extract_parameter_name(param)
                
                param_id = self._create_node(
                    'PARAMETER',
                    'parameter',
                    parameter_name=param_name
                )
                
                self._add_edge(func_id, param_id, 'has_param')
    def _enqueue_function_body(self, body: Dict, func_id: str):
        """将函数体加入队列"""
        if body is None:  # 添加None检查
            return
            
        if body.get('type') == 'BlockStatement':
            body_statements = body.get('body', [])
            for stmt in body_statements:
                if stmt is not None:  # 添加None检查
                    self._enqueue_child(stmt, func_id, 'function_body')
        else:
            self._enqueue_child(body, func_id, 'function_body')
    def _enqueue_child(self, child_node: Dict, parent_id: str, context: str):
        """将子节点加入队列"""
        if child_node is not None and child_node.get('type'):  # 添加None检查
            self.processing_queue.append({
                'ast_node': child_node,
                'parent_id': parent_id,
                'context': context
            })
    def _create_node(self, node_type: str, label: str, **attributes) -> str:
        """创建节点"""
        node_id = f"n{self.node_id_counter}"
        self.node_id_counter += 1
        
        node = {
            'id': node_id,
            'type': node_type,
            'label': label,
            **attributes
        }
        
        self.nodes.append(node)
        return node_id
    
    def _add_edge(self, source: str, target: str, edge_type: str):
        """添加边"""
        edge_exists = any(
            e['source'] == source and e['target'] == target and e['type'] == edge_type
            for e in self.edges
        )
        
        if not edge_exists:
            self.edges.append({
                'source': source,
                'target': target,
                'type': edge_type
            })
    
    def _extract_expression_name(self, expr: Dict) -> str:
        """提取表达式名称"""
        if expr is None:  # 添加None检查
            return 'unknown'
        
        expr_type = expr.get('type')
        
        if expr_type == 'Identifier':
            return expr.get('name', 'unknown')
        elif expr_type == 'MemberExpression':
            obj = expr.get('object', {})
            prop = expr.get('property', {})
            
            obj_name = self._extract_expression_name(obj) if obj is not None else 'unknown'
            prop_name = prop.get('name', 'unknown') if prop is not None else 'unknown'
            
            return f"{obj_name}.{prop_name}"
        elif expr_type == 'ThisExpression':
            return 'this'
        else:
            return expr_type.lower()
    def _extract_constructor_name(self, ast_node: Dict) -> str:
        """提取构造函数名"""
        callee = ast_node.get('callee', {})
        return self._extract_expression_name(callee)
    
    def _extract_parameter_name(self, param: Dict) -> str:
        """提取参数名"""
        if param.get('type') == 'Identifier':
            return param.get('name', 'unknown')
        elif param.get('type') == 'ObjectPattern':
            return '{...}'
        elif param.get('type') == 'ArrayPattern':
            return '[...]'
        else:
            return param.get('type', 'unknown')
    
    def _validate_and_generate_stats(self) -> Dict:
        """验证图结构并生成统计"""
        validation = self._validate_tree_structure()
        stats = self._generate_statistics()
        
        return {
            'nodes': self.nodes,
            'edges': self.edges,
            'validation': validation,
            'statistics': stats,
            'call_analysis': self._generate_call_analysis(),
            'summary': {
                'node_count': len(self.nodes),
                'edge_count': len(self.edges),
                'is_valid_tree': validation.get('is_valid_tree', False),
                'has_calls': any(e['type'] == 'chained_call' for e in self.edges)
            }
        }
    
    def _validate_tree_structure(self) -> Dict:
        """验证树结构"""
        validation = {
            'is_valid_tree': True,
            'errors': [],
            'warnings': []
        }
        
        # 检查入度
        in_degrees = defaultdict(int)
        for edge in self.edges:
            in_degrees[edge['target']] += 1
        
        # 找到根节点
        root_nodes = []
        for node in self.nodes:
            if in_degrees[node['id']] == 0:
                root_nodes.append(node)
        
        # 只有FILE类型应该是根节点
        actual_roots = [node for node in root_nodes if node['type'] == 'FILE']
        
        if len(actual_roots) != 1:
            validation['is_valid_tree'] = False
            validation['errors'].append(f"Expected 1 FILE root node, found {len(actual_roots)}")
        
        validation['root_nodes'] = actual_roots
        
        return validation
    
    def _generate_statistics(self) -> Dict:
        """生成统计信息"""
        stats = {
            'node_types': defaultdict(int),
            'edge_types': defaultdict(int),
            'function_count': 0,
            'call_count': 0,
            'class_count': 0,
            'chained_call_count': 0
        }
        
        for node in self.nodes:
            stats['node_types'][node['type']] += 1
            
            if node['type'] in ['FUNCTION_DEF', 'ANONYMOUS_FUNCTION']:
                stats['function_count'] += 1
            elif node['type'] == 'CALL':
                stats['call_count'] += 1
                if node.get('is_chained', False):
                    stats['chained_call_count'] += 1
            elif node['type'] == 'CLASS':
                stats['class_count'] += 1
        
        for edge in self.edges:
            stats['edge_types'][edge['type']] += 1
        
        return dict(stats)
    
    def _generate_call_analysis(self) -> Dict:
        """生成调用分析"""
        chained_calls = sum(1 for edge in self.edges if edge['type'] == 'chained_call')
        base_calls = len(set(edge['source'] for edge in self.edges if edge['type'] == 'chained_call'))
        
        return {
            'total_calls': chained_calls,
            'base_calls': base_calls,
            'call_details': {}
        }
    
    def save_pdcg(self, result: Dict, output_path: str):
        """保存PDCG数据"""
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(result, f, indent=2, ensure_ascii=False)
            print(f"PDCG分析结果已保存到: {output_path}")
        except Exception as e:
            print(f"保存失败: {e}")

    def print_single_file_summary(result: Dict):
        """打印单文件处理摘要"""
        print(f"\nPDCG分析结果摘要:")
        print(f"   节点数: {result['summary']['node_count']}")
        print(f"   边数: {result['summary']['edge_count']}")
        print(f"   链式调用数: {result['call_analysis']['total_calls']}")
        print(f"   基础调用数: {result['call_analysis']['base_calls']}")
        print(f"   是否为树: {result['summary']['is_valid_tree']}")
    
    def validate_directories(ast_root: str, pdcg_root: str) -> bool:
        """验证目录结构"""
        if not os.path.exists(ast_root):
            print(f"错误: AST根目录不存在: {ast_root}")
            return False
        
        # 创建PDCG根目录（如果不存在）
        try:
            os.makedirs(pdcg_root, exist_ok=True)
            print(f"PDCG根目录已准备: {pdcg_root}")
        except Exception as e:
            print(f"错误: 无法创建PDCG根目录: {e}")
            return False
        
        return True
class BatchPDCGProcessor:
    """批量PDCG处理器"""
    
    def __init__(self, ast_root: str, pdcg_root: str):
        self.ast_root = ast_root
        self.pdcg_root = pdcg_root
        self.processed_count = 0
        self.failed_count = 0
        self.skipped_count = 0
        self.start_time = None
        
        # 详细失败跟踪
        self.failed_files = []  # 存储失败文件的详细信息
        self.failure_categories = {
            'ast_load_error': [],
            'ast_parse_error': [],
            'pdcg_generation_error': [],
            'file_save_error': [],
            'permission_error': [],
            'unknown_error': []
        }
    def process_batch(self, data_type: str = 'both', limit: int = -1, 
                     target_packages: List[str] = None, skip_existing: bool = False,
                     generate_report: bool = True):
        """批量处理AST文件"""
        print(f"开始批量处理PDCG")
        print(f"  AST根目录: {self.ast_root}")
        print(f"  PDCG根目录: {self.pdcg_root}")
        print(f"  数据类型: {data_type}")
        print(f"  处理限制: {limit if limit > 0 else '无限制'}")
        print(f"  指定包: {target_packages if target_packages else '全部'}")
        print(f"  跳过已存在: {skip_existing}")
        print(f"  生成失败报告: {generate_report}")
        
        self.start_time = time.time()
        
        # 获取处理目标
        targets = self._get_processing_targets(data_type, target_packages, limit)
        
        if not targets:
            print("没有找到需要处理的AST文件")
            return
        
        print(f"找到 {len(targets)} 个AST文件需要处理")
        
        # 处理每个AST文件
        for i, (ast_file, pdcg_file) in enumerate(targets, 1):
            if skip_existing and os.path.exists(pdcg_file):
                print(f"[{i}/{len(targets)}] 跳过已存在: {os.path.basename(ast_file)}")
                self.skipped_count += 1
                continue
            
            self._process_single_ast(ast_file, pdcg_file, i, len(targets))
        
        # 输出总结
        self._print_batch_summary(len(targets))
        
        # 生成失败报告
        if generate_report and self.failed_files:
            self.generate_failure_report()
    def _get_processing_targets(self, data_type: str, target_packages: List[str], 
                               limit: int) -> List[tuple]:
        """获取处理目标列表"""
        targets = []
        
        # 确定要处理的数据类型目录
        data_dirs = []
        if data_type in ['both', 'benign']:
            benign_dir = os.path.join(self.ast_root, 'benign')
            if os.path.exists(benign_dir):
                data_dirs.append(('benign', benign_dir))
        
        if data_type in ['both', 'malicious']:
            malicious_dir = os.path.join(self.ast_root, 'malicious')
            if os.path.exists(malicious_dir):
                data_dirs.append(('malicious', malicious_dir))
        
        # 遍历每个数据类型目录
        for data_type_name, data_dir in data_dirs:
            package_dirs = [d for d in os.listdir(data_dir) 
                           if os.path.isdir(os.path.join(data_dir, d))]
            
            # 过滤指定包名
            if target_packages:
                package_dirs = [d for d in package_dirs if d in target_packages]
            
            # 遍历每个包目录
            for package_name in package_dirs:
                package_path = os.path.join(data_dir, package_name)
                ast_files = self._find_ast_files(package_path)
                
                for ast_file in ast_files:
                    # 计算对应的PDCG文件路径，将.ast.json替换为.pdcg.json
                    rel_path = os.path.relpath(ast_file, self.ast_root)
                    
                    # 替换文件后缀：.ast.json -> .pdcg.json
                    if rel_path.endswith('.ast.json'):
                        pdcg_rel_path = rel_path[:-9] + '.pdcg.json'  # 移除'.ast.json'，添加'.pdcg.json'
                    else:
                        # 如果不是.ast.json结尾，直接添加.pdcg.json
                        pdcg_rel_path = rel_path + '.pdcg.json'
                    
                    pdcg_file = os.path.join(self.pdcg_root, pdcg_rel_path)
                    
                    targets.append((ast_file, pdcg_file))
                    
                    # 检查限制
                    if limit > 0 and len(targets) >= limit:
                        return targets
        
        return targets
    def _find_ast_files(self, package_path: str) -> List[str]:
        """在包目录中查找所有AST文件"""
        ast_files = []
        
        for root, dirs, files in os.walk(package_path):
            for file in files:
                if file.endswith('.ast.json'):
                    ast_files.append(os.path.join(root, file))
        
        return sorted(ast_files)
    
    def _process_single_ast(self, ast_file: str, pdcg_file: str, current: int, total: int):
        """处理单个AST文件"""
        package_name = self._extract_package_name(ast_file)
        file_name = os.path.basename(ast_file)
        
        try:
            print(f"[{current}/{total}] 处理: {package_name}/{file_name}")
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(pdcg_file), exist_ok=True)
            
            # 创建分析器 - 修改这里，不传递参数
            analyzer = PDCGAnalyzer()
            
            # 尝试分析
            result = analyzer.analyze_pdcg_from_ast(ast_file)
            
            if 'error' in result:
                # PDCG生成失败
                error_msg = result['error']
                print(f"  ❌ 分析失败: {error_msg}")
                
                self._record_failure(
                    ast_file=ast_file,
                    pdcg_file=pdcg_file,
                    package_name=package_name,
                    file_name=file_name,
                    error_type='pdcg_generation_error',
                    error_message=error_msg,
                    stage='PDCG Generation'
                )
                return
            
            # 尝试保存结果
            try:
                analyzer.save_pdcg(result, pdcg_file)
            except PermissionError as e:
                self._record_failure(
                    ast_file=ast_file,
                    pdcg_file=pdcg_file,
                    package_name=package_name,
                    file_name=file_name,
                    error_type='permission_error',
                    error_message=f"权限错误: {str(e)}",
                    stage='File Save'
                )
                return
            except OSError as e:
                self._record_failure(
                    ast_file=ast_file,
                    pdcg_file=pdcg_file,
                    package_name=package_name,
                    file_name=file_name,
                    error_type='file_save_error',
                    error_message=f"文件保存错误: {str(e)}",
                    stage='File Save'
                )
                return
            
            # 成功处理
            node_count = result['summary']['node_count']
            edge_count = result['summary']['edge_count']
            print(f"  ✅ 成功: {node_count}个节点, {edge_count}条边")
            
            self.processed_count += 1
            
        except FileNotFoundError as e:
            self._record_failure(
                ast_file=ast_file,
                pdcg_file=pdcg_file,
                package_name=package_name,
                file_name=file_name,
                error_type='ast_load_error',
                error_message=f"AST文件不存在: {str(e)}",
                stage='AST Load'
            )
            
        except json.JSONDecodeError as e:
            self._record_failure(
                ast_file=ast_file,
                pdcg_file=pdcg_file,
                package_name=package_name,
                file_name=file_name,
                error_type='ast_parse_error',
                error_message=f"AST解析错误: {str(e)}",
                stage='AST Parse'
            )
            
        except Exception as e:
            self._record_failure(
                ast_file=ast_file,
                pdcg_file=pdcg_file,
                package_name=package_name,
                file_name=file_name,
                error_type='unknown_error',
                error_message=f"未知错误: {str(e)}",
                stage='Unknown'
            )
    def _record_failure(self, ast_file: str, pdcg_file: str, package_name: str, 
                       file_name: str, error_type: str, error_message: str, stage: str):
        """记录失败详情"""
        failure_info = {
            'package_name': package_name,
            'file_name': file_name,
            'ast_file': ast_file,
            'pdcg_file': pdcg_file,
            'error_type': error_type,
            'error_message': error_message,
            'failure_stage': stage,
            'timestamp': time.time()
        }
        
        # 添加到总失败列表
        self.failed_files.append(failure_info)
        
        # 按类别分类
        if error_type in self.failure_categories:
            self.failure_categories[error_type].append(failure_info)
        
        self.failed_count += 1
    def generate_failure_report(self, output_dir: str = None):
        """生成详细的失败报告"""
        if not self.failed_files:
            print("没有失败的文件，无需生成失败报告")
            return
        
        # 确定输出目录
        if output_dir is None:
            output_dir = os.path.dirname(self.pdcg_root)
        
        os.makedirs(output_dir, exist_ok=True)
        
        # 生成文本报告
        report_file = os.path.join(output_dir, "pdcg_failure_report.txt")
        self._generate_text_report(report_file)
        
        # 生成CSV报告
        csv_file = os.path.join(output_dir, "pdcg_failure_details.csv")
        self._generate_csv_report(csv_file)
        
        print(f"\n失败报告已生成:")
        print(f"  详细报告: {report_file}")
        print(f"  CSV文件: {csv_file}")
    
    def _generate_text_report(self, report_file: str):
        """生成文本格式的失败报告"""
        with open(report_file, 'w', encoding='utf-8') as f:
            f.write("="*80 + "\n")
            f.write("PDCG生成失败详细报告\n")
            f.write("="*80 + "\n")
            f.write(f"生成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"总失败文件数: {len(self.failed_files)}\n\n")
            
            # 按错误类型分类统计
            f.write("1. 失败类型统计\n")
            f.write("-"*40 + "\n")
            for error_type, failures in self.failure_categories.items():
                if failures:
                    f.write(f"{error_type}: {len(failures)} 个文件\n")
            f.write("\n")
            
            # 按包名分组统计
            package_failures = {}
            for failure in self.failed_files:
                pkg = failure['package_name']
                if pkg not in package_failures:
                    package_failures[pkg] = []
                package_failures[pkg].append(failure)
            
            f.write("2. 按包名分组的失败统计\n")
            f.write("-"*40 + "\n")
            for pkg_name, pkg_failures in sorted(package_failures.items()):
                f.write(f"{pkg_name}: {len(pkg_failures)} 个文件失败\n")
            f.write("\n")
            
            # 详细失败信息
            f.write("3. 详细失败信息\n")
            f.write("-"*40 + "\n")
            
            for i, failure in enumerate(self.failed_files, 1):
                f.write(f"\n失败 #{i}:\n")
                f.write(f"  包名: {failure['package_name']}\n")
                f.write(f"  文件: {failure['file_name']}\n")
                f.write(f"  AST路径: {failure['ast_file']}\n")
                f.write(f"  错误类型: {failure['error_type']}\n")
                f.write(f"  失败阶段: {failure['failure_stage']}\n")
                f.write(f"  错误信息: {failure['error_message']}\n")
                f.write(f"  时间戳: {time.strftime('%H:%M:%S', time.localtime(failure['timestamp']))}\n")
            
            # 错误类型详细分析
            f.write("\n\n4. 错误类型详细分析\n")
            f.write("="*40 + "\n")
            
            for error_type, failures in self.failure_categories.items():
                if not failures:
                    continue
                    
                f.write(f"\n{error_type.upper()} ({len(failures)} 个文件):\n")
                f.write("-"*30 + "\n")
                
                # 按包名分组显示
                type_packages = {}
                for failure in failures:
                    pkg = failure['package_name']
                    if pkg not in type_packages:
                        type_packages[pkg] = []
                    type_packages[pkg].append(failure['file_name'])
                
                for pkg_name, file_names in sorted(type_packages.items()):
                    f.write(f"  {pkg_name}:\n")
                    for file_name in sorted(file_names):
                        f.write(f"    - {file_name}\n")
            
            f.write("\n" + "="*80 + "\n")
            f.write("报告结束\n")
            f.write("="*80 + "\n")
    
    def _generate_csv_report(self, csv_file: str):
        """生成CSV格式的失败报告"""
        import pandas as pd
        
        if not self.failed_files:
            return
        
        # 转换为DataFrame
        df = pd.DataFrame(self.failed_files)
        
        # 添加时间戳转换
        df['formatted_time'] = df['timestamp'].apply(
            lambda x: time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(x))
        )
        
        # 重新排列列
        columns_order = [
            'package_name', 'file_name', 'error_type', 'failure_stage',
            'error_message', 'formatted_time', 'ast_file', 'pdcg_file'
        ]
        
        df = df[columns_order]
        
        # 保存CSV
        df.to_csv(csv_file, index=False, encoding='utf-8')
    def _extract_package_name(self, ast_file: str) -> str:
        """从AST文件路径提取包名"""
        # 从路径中提取包名
        # E:\AST\malicious\package_name\... -> package_name
        parts = ast_file.replace('\\', '/').split('/')
        
        for i, part in enumerate(parts):
            if part in ['benign', 'malicious']:
                if i + 1 < len(parts):
                    return parts[i + 1]
        
        return 'unknown_package'
    
    def _print_batch_summary(self, total_targets: int):
        """打印批量处理总结"""
        elapsed_time = time.time() - self.start_time
        
        print(f"\n" + "="*80)
        print(f"批量PDCG处理完成!")
        print(f"="*80)
        print(f"  总目标数: {total_targets}")
        print(f"  成功处理: {self.processed_count}")
        print(f"  处理失败: {self.failed_count}")
        print(f"  跳过文件: {self.skipped_count}")
        print(f"  处理时间: {elapsed_time:.2f} 秒")
        
        if self.processed_count > 0:
            avg_time = elapsed_time / self.processed_count
            print(f"  平均时间: {avg_time:.2f} 秒/文件")
        
        success_rate = (self.processed_count / total_targets * 100) if total_targets > 0 else 0
        print(f"  成功率: {success_rate:.1f}%")
        
        # 失败详细统计
        if self.failed_files:
            print(f"\n失败详细统计:")
            print("-"*40)
            
            for error_type, failures in self.failure_categories.items():
                if failures:
                    print(f"  {error_type}: {len(failures)} 个文件")
            
            # 显示失败最多的包
            package_failures = {}
            for failure in self.failed_files:
                pkg = failure['package_name']
                package_failures[pkg] = package_failures.get(pkg, 0) + 1
            
            if package_failures:
                print(f"\n失败最多的包 (Top 10):")
                sorted_packages = sorted(package_failures.items(), key=lambda x: x[1], reverse=True)
                for pkg_name, count in sorted_packages[:10]:
                    print(f"    {pkg_name}: {count} 个文件失败")
            
            # 显示典型失败示例
            print(f"\n典型失败示例:")
            for error_type, failures in self.failure_categories.items():
                if failures:
                    example = failures[0]
                    print(f"  {error_type}:")
                    print(f"    包: {example['package_name']}")
                    print(f"    文件: {example['file_name']}")
                    print(f"    错误: {example['error_message'][:100]}...")
                    break
        
        print(f"="*80)
def print_single_file_summary(result: Dict):
    """打印单文件处理摘要"""
    print(f"\nPDCG分析结果摘要:")
    print(f"   节点数: {result['summary']['node_count']}")
    print(f"   边数: {result['summary']['edge_count']}")
    print(f"   链式调用数: {result['call_analysis']['total_calls']}")
    print(f"   基础调用数: {result['call_analysis']['base_calls']}")
    print(f"   是否为树: {result['summary']['is_valid_tree']}")

def validate_directories(ast_root: str, pdcg_root: str) -> bool:
    """验证目录结构"""
    if not os.path.exists(ast_root):
        print(f"错误: AST根目录不存在: {ast_root}")
        return False
    
    # 创建PDCG根目录（如果不存在）
    try:
        os.makedirs(pdcg_root, exist_ok=True)
        print(f"PDCG根目录已准备: {pdcg_root}")
    except Exception as e:
        print(f"错误: 无法创建PDCG根目录: {e}")
        return False
    
    return True
def main():
    """主函数 - 支持批量处理"""
    import argparse
    
    parser = argparse.ArgumentParser(description='PDCG分析器 - 支持单文件和批量处理')
    parser.add_argument('--mode', choices=['single', 'batch'], default='single', help='处理模式')
    parser.add_argument('--input', type=str, help='单文件模式：AST文件路径；批量模式：AST根目录路径')
    parser.add_argument('--output', type=str, help='单文件模式：输出文件路径；批量模式：PDCG根目录路径')
    parser.add_argument('--type', choices=['benign', 'malicious', 'both'], default='both', help='批量模式：处理数据类型')
    parser.add_argument('--limit', type=int, default=-1, help='批量模式：限制处理包数量（-1表示全部）')
    parser.add_argument('--packages', nargs='*', help='批量模式：指定包名列表')
    parser.add_argument('--skip-existing', action='store_true', help='批量模式：跳过已存在的PDCG文件')
    parser.add_argument('--no-report', action='store_true', help='批量模式：不生成失败报告')
    
    args = parser.parse_args()
    
    print("启动PDCG分析器")
    
    if args.mode == 'single':
        # 单文件处理模式（保持不变）
        ast_file = args.input or r"H:\NEW_AST_P\malicious\1_postmessage-api-external-3.25.12\package\build.js.ast.json"
        output_file = args.output or r"E:\PythonProject\malicious_npm_detector_1.0.1\test_pdcg.json"
        
        if not os.path.exists(ast_file):
            print(f"AST文件不存在: {ast_file}")
            return
        
        analyzer = PDCGAnalyzer()
        result = analyzer.analyze_pdcg_from_ast(ast_file)
        
        if 'error' in result:
            print(f"分析失败: {result['error']}")
            return
        
        analyzer.save_pdcg(result, output_file)
        print_single_file_summary(result)
        
    elif args.mode == 'batch':
        # 批量处理模式
        ast_root = args.input or r"G:\extracted_packages_ast"
        pdcg_root = args.output or r"G:\extracted_packages_PDCG"
        
        if not os.path.exists(ast_root):
            print(f"AST根目录不存在: {ast_root}")
            return
        
        batch_processor = BatchPDCGProcessor(ast_root, pdcg_root)
        batch_processor.process_batch(
            data_type=args.type,
            limit=args.limit,
            target_packages=args.packages,
            skip_existing=args.skip_existing,
            generate_report=not args.no_report
        )
    
    print(f"\nPDCG分析完成!")
if __name__ == "__main__":
    main()